#!/usr/bin/python3
# coding=utf-8
#
# Copyright (C) 2025. Huawei Technologies Co., Ltd. All rights reserved.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
# ===============================================================================

import numpy as np
import os

# 测试用例编号
id = 9

def custom_maskfill(input_x, mask, value):
    """自定义MaskFill实现,支持广播"""
    output = np.copy(input_x)
    broadcasted_mask = np.broadcast_to(mask, input_x.shape)
    output[broadcasted_mask] = value 
    return output


def gen_golden_data_maskfill(case_id):
    # 创建目录
    input_dir = "./input"
    output_dir = "./output"
    os.makedirs(input_dir, exist_ok=True)
    os.makedirs(output_dir, exist_ok=True)
    
    # 初始化参数
    input_shape_x = []
    input_shape_mask = []
    input_dtype = np.float32
    input_range = (0, 1)
    mask_true_ratio = 0.5
    fill_value = 0
    
    if case_id == 0:
        # case 0: 3D维度 + float16
        input_shape_x = [512, 512, 512]      
        input_shape_mask = [1, 1, 512] 
        input_dtype = np.float16      
        input_range = (-3000, 3000)   
        mask_true_ratio = 0.4         
        fill_value = -1               
    elif case_id == 1:
        # case 1: 2D维度 + float32
        input_shape_x = [4096, 2048]        
        input_shape_mask = [4096, 1]         
        input_dtype = np.float32             
        input_range = (-1.23456789, 1.23456789)   
        mask_true_ratio = 0.6                  
        fill_value = -1.23456789               
    elif case_id == 2:
        # case 2: 5维嵌套维度 + int32
        input_shape_x = [1, 1, 1, 1, 100000]  
        input_shape_mask = [1, 1, 1, 1, 1]     
        input_dtype = np.int32                 
        input_range = (-2147483647, 2147483647)  
        mask_true_ratio = 0.2                   
        fill_value = 2147483647                
    elif case_id == 3:
        # case 3: 1D维度 + uint8
        input_shape_x = [65536]                
        input_shape_mask = [1]                  
        input_dtype = np.uint8                 
        input_range = (0, 255)                 
        mask_true_ratio = 0.5                  
        fill_value = 255                       
    elif case_id == 4:
        # case 4: 1 + bool
        input_shape_x = [1]              
        input_shape_mask = [1]           
        input_dtype = np.bool_                 
        input_range = (0, 1)                   
        mask_true_ratio = 0.3                  
        fill_value = False                     
    elif case_id == 5:
        # case 5: 6维高维 + int64
        input_shape_x = [16, 1, 1, 1, 1, 65535]  
        input_shape_mask = [1, 1, 1, 1, 1, 65535]  
        input_dtype = np.int64                   
        input_range = (-9223372036854775808, 0)  
        mask_true_ratio = 0.7                    
        fill_value = -9223372036854775808        
    elif case_id == 6:
        # case 6: 3D维度 + bfloat16（用float32模拟）
        input_shape_x = [2048, 2048, 2]        
        input_shape_mask = [2048, 1, 1]           
        input_dtype = np.float32  # 实际场景替换为bfloat16
        input_range = (0, 3.1415)               
        mask_true_ratio = 0.1                   
        fill_value = 3.1415                     
    elif case_id == 7:
        # case 7: 4D非对称广播 + double
        input_shape_x = [1, 1638, 1, 1638]    
        input_shape_mask = [1, 1, 1, 1638]     
        input_dtype = np.double                 
        input_range = (0, 1.7976931348623157e+308)  
        mask_true_ratio = 0.8                       
        fill_value = 1.7976931348623157e+308        
    elif case_id == 8:
        # case 8: 混合空维度 + int16
        input_shape_x = [32, 1, 8192]           
        input_shape_mask = [32, 1, 1]           
        input_dtype = np.int16                  
        input_range = (-32768, 0)               
        mask_true_ratio = 0.9                   
        fill_value = -32768                     
    elif case_id == 9:
        # case 9: 7维超长维度 + uint64
        input_shape_x = [1, 1, 1, 1, 1, 1, 262144]  
        input_shape_mask = [1, 1, 1, 1, 1, 1, 1]      
        input_dtype = np.uint64                       
        input_range = (0, 18446744073709551615)       
        mask_true_ratio = 0.05                       
        fill_value = 18446744073709551615              
    else:
        raise ValueError(f"Not sopported.")
    
    # 生成输入数据
    input_x = np.random.uniform(
        low=input_range[0], 
        high=input_range[1], 
        size=input_shape_x
    ).astype(input_dtype)
    
    # 生成掩码
    if len(input_shape_mask) == 0:
        mask = np.array(np.random.choice(a=[False, True], p=[1 - mask_true_ratio, mask_true_ratio]), dtype=np.bool_)
    else:
        mask = np.random.choice(
            a=[False, True], 
            size=input_shape_mask, 
            p=[1 - mask_true_ratio, mask_true_ratio]
        ).astype(np.bool_)
    
    # 生成填充值
    fill_value_casted = np.array(fill_value, dtype=input_dtype)
    
    # 计算金标准
    golden = custom_maskfill(input_x, mask, fill_value_casted)
    
    # 保存数据
    input_x.tofile(os.path.join(input_dir, "input_x.bin"))
    mask.tofile(os.path.join(input_dir, "mask.bin"))
    fill_value_casted.tofile(os.path.join(input_dir, "fill_value.bin"))
    golden.tofile(os.path.join(output_dir, "golden.bin"))
    print(f"Case {case_id} 数据生成完成，路径：{input_dir} 和 {output_dir}")

if __name__ == "__main__":
    gen_golden_data_maskfill(case_id=id)